
#import "UIApplicationAdditions.h"
#include <mach-o/arch.h>
#include <mach-o/fat.h>
#include <mach-o/loader.h>

//
// blob
//
typedef struct _blob
{
	uint32_t magic;
	uint32_t length;
}blob;

//
// blob index
//
typedef struct _cs_blob_index
{
	uint32_t type;
	uint32_t offset;
}cs_blob_index;

//
// super
//
typedef struct _cs_super_block
{
	uint32_t magic;
	uint32_t length;
	uint32_t count;
	cs_blob_index index[];
}cs_super_block;

#define CSMAGIC_EMBEDDED_SIGNATURE											0xfade0cc0
#define CSMAGIC_ENTITLEMENTS												0xfade7171
#define CSSLOT_ENTITLEMENTS													0x00000005

//
// swap routine
//
typedef uint32_t (*fnSwap32)(uint32_t);

//
// noop
//
static uint32_t swap32Noop(uint32_t value)
{
	return value;
}

//
// swap
//
static uint32_t swap32(uint32_t value)
{
	return OSSwapInt32(value);
}

@implementation UIApplication(Additions)

//
// parse mach-o
//
- (NSDictionary*)entitlementsFromMachOBuffer:(uint8_t const*)machOBuffer length:(uint32_t)length
{
	uint32_t magic															= *(uint32_t const*)machOBuffer;
	struct mach_header const* header										= 0;
	uint32_t headerSize														= sizeof(struct mach_header);
	struct fat_header const* fatHeader										= 0;
	fnSwap32 swap32Routine													= &swap32Noop;
	switch(magic)
	{
	case MH_CIGAM:
		swap32Routine														= &swap32;
		
	case MH_MAGIC:
		header																= (struct mach_header const*)machOBuffer;
		break;
		
	case MH_CIGAM_64:
		swap32Routine														= &swap32;
		
	case MH_MAGIC_64:
		headerSize															= sizeof(struct mach_header_64);
		header																= (struct mach_header const*)machOBuffer;
		break;
		
	case FAT_CIGAM:
	case FAT_MAGIC:
		{
			fatHeader														= (struct fat_header const*)machOBuffer;
			uint32_t fatCount												= OSSwapBigToHostInt32(fatHeader->nfat_arch);
			if(!fatCount)
				return nil;
			
			struct fat_arch const* fatArch									= (struct fat_arch const*)(machOBuffer + sizeof(struct fat_header));
			machOBuffer														+= OSSwapBigToHostInt32(fatArch->offset);
			length															= OSSwapBigToHostInt32(fatArch->size);
			return [self entitlementsFromMachOBuffer:machOBuffer length:length];
		}
		break;
			
	default:
		return nil;
		break;
	}
	
	struct load_command const* loadCommand									= (struct load_command const*)((uint8_t const*)header + headerSize);
	uint32_t commandCount													= swap32Routine(header->ncmds);
	for(uint32_t i = 0; i < commandCount; i ++)
	{
		uint32_t commandSize												= swap32Routine(loadCommand->cmdsize);
		uint32_t commandType												= swap32Routine(loadCommand->cmd);
		if(commandType == LC_CODE_SIGNATURE)
		{
			struct linkedit_data_command const* codeSignatureCommand		= (struct linkedit_data_command const*)loadCommand;
			uint32_t dataOffset												= swap32Routine(codeSignatureCommand->dataoff);
			uint8_t const* codeSignatureBuffer								= machOBuffer + dataOffset;
			cs_super_block const* superBlock								= (cs_super_block const*)codeSignatureBuffer;
			if(OSSwapBigToHostInt32(superBlock->magic) == CSMAGIC_EMBEDDED_SIGNATURE)
			{
				uint32_t count												= OSSwapBigToHostInt32(superBlock->count);
				for(uint32_t j = 0; j < count; j ++)
				{
					if(OSSwapBigToHostInt32(superBlock->index[j].type) != CSSLOT_ENTITLEMENTS)
						continue;
					
					uint32_t offset											= OSSwapBigToHostInt32(superBlock->index[j].offset);
					blob const* blobHeader									= (blob const*)(codeSignatureBuffer + offset);
					if(OSSwapBigToHostInt32(blobHeader->magic) != CSMAGIC_ENTITLEMENTS)
						continue;
					
					uint32_t dataLength										= OSSwapBigToHostInt32(blobHeader->length) - sizeof(blob);
					NSData* data											= [NSData dataWithBytes:blobHeader + 1 length:dataLength];
					NSError* error											= nil;
					NSPropertyListReadOptions format						= NSPropertyListXMLFormat_v1_0;
					return [NSPropertyListSerialization propertyListWithData:data options:0 format:&format error:&error];
				}
			}
		}
		
		loadCommand															= (struct load_command const*)((uint8_t const*)loadCommand + commandSize);
	}
	
	return nil;
}

//
// get entitlements
//
- (NSDictionary*)entitlements
{
	NSData* mainExecutableData												= [NSData dataWithContentsOfURL:[[NSBundle mainBundle] executableURL]];
	return [self entitlementsFromMachOBuffer:(uint8_t const*)mainExecutableData.bytes length:mainExecutableData.length];
}

@end
